import time
import datetime
from typing import List

import joblib
import pandas as pd

from dynamic_compute import compute
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed


def parallel_compute(params_lst: List[dict]) -> []:
    pool = ProcessPoolExecutor(max_workers=5)
    tasks = [pool.submit(compute, seq_id, params) for seq_id, params in enumerate(params_lst)]
    results = [t.result() for t in as_completed(tasks)]
    return sorted(results, key=lambda x: x["seq_id"])


if __name__ == "__main__":
    calc_begin_time = time.perf_counter()
    result = {
        "h_m": {},
        "c_e": {},
        "h_e": {}
    }
    v_0 = 1
    v_1 = 8
    v_2 = 12
    v_3 = 15
    T = 8

    # h_m
    params_lst = [{"_T": T, "_v_0": v_0, "_v_1": v_1, "_v_2": v_2, "_v_3": v_3, "_b": 20, "_c_m": 2, "_h_m": h_m, "_c_e": 10, "_h_e": 1}
                  for h_m in [0.1, 0.15, 0.2, 0.25, 0.3]]
    res_h_m = parallel_compute(params_lst)
    for res in res_h_m:
        result["h_m"].update({res["seq_id"]: res})

    # c_e
    params_lst = [{"_T": T, "_v_0": v_0, "_v_1": v_1, "_v_2": v_2, "_v_3": v_3, "_b": 20, "_c_m": 2, "_h_m": 0.2, "_c_e": c_e, "_h_e": 1}
                  for c_e in [6, 8, 10, 12, 14]]
    res_c_e = parallel_compute(params_lst)
    for res in res_c_e:
        result["c_e"].update({res["seq_id"]: res})

    # h_e
    params_lst = [{"_T": T, "_v_0": v_0, "_v_1": v_1, "_v_2": v_2, "_v_3": v_3, "_b": 20, "_c_m": 2, "_h_m": 0.2, "_c_e": 10, "_h_e": h_e}
                  for h_e in [0.6, 0.8, 1, 1.2, 1.4]]
    res_h_e = parallel_compute(params_lst)
    for res in res_h_e:
        result["h_e"].update({res["seq_id"]: res})

    calc_end_time = time.perf_counter()
    print(f"计算完成, 耗时 {calc_end_time - calc_begin_time: .2f} 秒")
    # output

    joblib.dump(result, f"results/{datetime.datetime.now().strftime('%Y-%m-%d %H-%M-%S')}.pkl")
    writer = pd.ExcelWriter(f"results/{datetime.datetime.now().strftime('%Y-%m-%d %H-%M-%S')}.xlsx")
    data_h_m = pd.DataFrame(result["h_m"])
    data_c_e = pd.DataFrame(result["c_e"])
    data_h_e = pd.DataFrame(result["h_e"])
    data_h_m.to_excel(writer, sheet_name="h_m")
    data_c_e.to_excel(writer, sheet_name="c_e")
    data_h_e.to_excel(writer, sheet_name="h_e")
    writer.close()
